/*
 * @Author: your name
 * @Date: 2021-10-11 10:01:51
<<<<<<< HEAD
 * @LastEditTime: 2021-10-17 11:23:29
=======
 * @LastEditTime: 2021-10-13 11:46:58
>>>>>>> fc51075a61ea3d0ec07cbba77045e13b1c51b3a5
 * @LastEditors: Please set LastEditors
 * @Description: In User Settings Edit
 * @FilePath: \test_algorithm\webrtcdemo.cpp
 */

#include <iostream>

#ifdef _WIN32
#include <Windows.h>
#pragma comment(lib, "d3d11.lib")
#endif

#define WEBRTC_WIN

#include "modules/desktop_capture/desktop_capturer.h"
#include "modules/desktop_capture/cropping_window_capturer.h"
#include "modules/desktop_capture/desktop_capture_options.h"
#include "modules/desktop_capture/screen_capturer_helper.h"
#include "modules/desktop_capture/desktop_frame.h"

#include <iostream>
#include "png.h"
/******************************图片数据*********************************/
typedef struct _pic_data pic_data;
struct _pic_data
{
    int width, height; /* 尺寸 */
    int bit_depth;     /* 位深 */
    int flag;          /* 一个标志，表示是否有alpha通道 */

    unsigned char **rgba; /* 图片数组 */
};
/**********************************************************************/
#define PNG_BYTES_TO_CHECK 4
#define HAVE_ALPHA 1
#define NO_ALPHA 0

int write_png_file(char *file_name, pic_data *graph)
/* 功能：将LCUI_Graph结构中的数据写入至png文件 */
{
    int j, i, temp, pos;
    png_byte color_type;

    png_structp png_ptr;
    png_infop info_ptr;
    png_bytep *row_pointers;
    /* create file */
    FILE *fp = fopen(file_name, "wb");
    if (!fp)
    {
        printf("[write_png_file] File %s could not be opened for writing", file_name);
        return -1;
    }

    /* initialize stuff */
    png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

    if (!png_ptr)
    {
        printf("[write_png_file] png_create_write_struct failed");
        return -1;
    }
    info_ptr = png_create_info_struct(png_ptr);
    if (!info_ptr)
    {
        printf("[write_png_file] png_create_info_struct failed");
        return -1;
    }
    
    if (setjmp(png_jmpbuf(png_ptr)))
    {
        printf("[write_png_file] Error during init_io");
        return -1;
    }
    png_init_io(png_ptr, fp);

    /* write header */
    if (setjmp(png_jmpbuf(png_ptr)))
    {
        printf("[write_png_file] Error during writing header");
        return -1;
    }
    /* 判断要写入至文件的图片数据是否有透明度，来选择色彩类型 */
    if (graph->flag == HAVE_ALPHA)
        color_type = PNG_COLOR_TYPE_RGB_ALPHA;
    else
        color_type = PNG_COLOR_TYPE_RGB;

    png_set_IHDR(png_ptr, info_ptr, graph->width, graph->height,
                 graph->bit_depth, color_type, PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    png_write_info(png_ptr, info_ptr);

    /* write bytes */
    if (setjmp(png_jmpbuf(png_ptr)))
    {
        printf("[write_png_file] Error during writing bytes");
        return -1;
    }
    if (graph->flag == HAVE_ALPHA)
        temp = (4 * graph->width);
    else
        temp = (3 * graph->width);

    pos = 0;
    row_pointers = (png_bytep *)malloc(graph->height * sizeof(png_bytep));
    for (i = 0; i < graph->height; i++)
    {
        row_pointers[i] = ((png_bytep )malloc(temp)); 
        memcpy(row_pointers[i],*graph->rgba + i*temp,temp);
    }
    png_write_image(png_ptr, row_pointers);

    /* end write */
    if (setjmp(png_jmpbuf(png_ptr)))
    {
        printf("[write_png_file] Error during end of write");
        return -1;
    }
    png_write_end(png_ptr, NULL);

    /* cleanup heap allocation */
    for (j = 0; j < graph->height; j++)
        free(row_pointers[j]);
    free(row_pointers);

    fclose(fp);
    return 0;
}

class CaptureCallBack : public webrtc::DesktopCapturer::Callback
{
    virtual void OnCaptureResult(webrtc::DesktopCapturer::Result ret,
        std::unique_ptr<webrtc::DesktopFrame> frame){
        if(ret == webrtc::DesktopCapturer::Result::SUCCESS){
            std::cout
                << "capture frame "
                << frame.get()->size().width() << " "
                << frame.get()->size().height() << "\r\n";
                                 std::unique_ptr<webrtc::DesktopFrame> frame)
    {
        if (ret == webrtc::DesktopCapturer::Result::SUCCESS)
        {
            std::cout << "capture frame " << frame.get()->size().width() << " "
                      << frame.get()->size().height() << "\r\n " << frame.get()->capture_time_ms() << " " << frame.get()->stride() << "\r\n";
            pic_data pic;
            pic.width = frame.get()->size().width();
            pic.height = frame.get()->size().height();
            auto data = frame.get()->data();
            pic.rgba = (uint8_t**)(&(data));
            pic.bit_depth = 8;
            pic.flag = 1;

            write_png_file("test.png",&pic);
            exit(0);
        }
    }
};

int main(int argc, char *argv[])
{
    auto screen_capture_ =
        std::move(webrtc::DesktopCapturer::CreateScreenCapturer(webrtc::DesktopCaptureOptions::CreateDefault()));
    screen_capture_->Start(new CaptureCallBack());
    while (true)
    {
        screen_capture_->CaptureFrame();
        Sleep(100);
    }
    return 0;
}
